#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <cmath>

#include "armadillo"
#include "matplotlibcpp.h"

// 主要参考 https://yunchengjiang.blog.csdn.net/article/details/103458473?spm=1001.2014.3001.5502
using namespace std;
using namespace arma;
namespace plt = matplotlibcpp;

typedef struct {
    double delta;
    uword index;
    double e;
    double th_e;
} lqrResult_t;

typedef struct {
    uword index;
    double latError;
} minLatError_idx_t;

const double deg2rad = M_PI/180;
const double rad2deg = 180/M_PI;
const double toKMperH = 3.6;
const double tompers = 1/3.6;
const char fileName[] = "/home/shukebeita/data/data.txt";

inline double normFunc(double posRefX, double posRefY, double posX, double posY){
    return sqrt(pow(posRefX - posX, 2) + pow(posRefY - posY, 2));
}

//读取路径点
mat readPosPoints(vector<double> &XRef, vector<double> &YRef){
    ifstream readFile;
    readFile.open(fileName);

    if(!readFile.is_open()){
        cout << "file failed !!!" << endl;
        exit(0); //表示程序正常退出
    }

    string dataBuff;
    while(getline(readFile, dataBuff)){
        string tempData;
        vector<string> tempStr;
        for(unsigned long i = 0; i <= dataBuff.length(); i ++){
            if(dataBuff[i] == ',' || dataBuff[i] == '\0'){
                tempStr.push_back(tempData);
                tempData.clear();
            }else{
                tempData += dataBuff[i];
            }
        }

        XRef.push_back(stod(tempStr[0]));
        YRef.push_back(stod(tempStr[1]));
    }

    readFile.close();

    // 注意：vector<double> 类型可以直接赋给 vec 反之，不行！！！
    vec XRefTemp = XRef;
    vec YRefTemp = YRef;
    mat MRes;
    MRes = join_horiz(XRefTemp, YRefTemp); //相当于MATLAB中的 A = [B, C];

    /*
    // 画出曲线，这里作为验证
    vector<double> X;
    vector<double> Y;
    for(uword i = 0; i < XRefTemp.n_elem; i++){
        X.push_back(XRefTemp[i]);
        Y.push_back(YRefTemp[i]);
    }
    plt::plot(X, Y, "r-");
    plt::show();
    */

    return MRes;
}

vec toPiVec(vec angle){
    vec res(angle.n_elem);
    for(uword i = 0; i < angle.n_elem; i++){
        if(angle[i] > M_PI)
            res[i] = angle[i] - 2*M_PI;
        else if(angle[i] < -M_PI)
            res[i] = angle[i] + 2*M_PI;
        else
            res[i] = angle[i];
    }
    return res;
}

double toPi(double angle){
    double res;
    if(angle > M_PI)
        res = angle - 2*M_PI;
    else if(angle < -M_PI)
        res = angle + 2*M_PI;
    else
        res = angle;

    return res;
}

minLatError_idx_t findMinLatError_idx(vec states, mat refData){
    uword lenOfPosRef = refData.n_rows;

    vec dist(lenOfPosRef);
    for(uword i = 0; i < lenOfPosRef; i++){
        dist(i) = normFunc(refData(i, 0), refData(i, 1), states(0), states(1));
    }
    double minLatError_temp = dist.min(); // 找到当前位置到路径点的最小距离
    uword idx = dist.index_min(); // 找到距离当前位置最近的一个参考轨迹点的序号

    minLatError_idx_t tempMinLatError_idx;
    tempMinLatError_idx.index = idx;

    if(states(1) < refData(idx, 1))
        tempMinLatError_idx.latError = -minLatError_temp;
    else
        tempMinLatError_idx.latError = minLatError_temp;

    return tempMinLatError_idx;
}

rowvec dlqr(mat A, mat B, mat Q, mat R){
    // first, solve the ricatti equation
    mat X = Q;
    mat Xnext = Q;

    int maxIter = 200;
    double epsilon = 0.01;
    for(int i = 0; i < maxIter; i++){
        //Xnext = A.t()*X*A - A.t()*X*B* inv(R+B.t()*X*B) *B.t()*X*A + Q;
        Xnext = A.t()*X*A - A.t()*X*B*solve(R+B.t()*X*B, B.t())*X*A + Q;
        if(abs(Xnext - X).max() <= epsilon){
            //X = Xnext;
            break;
        }
        X = Xnext;
    }

    //compute the LQR gain
    rowvec Ktemp = solve(B.t()*X*B + R, B.t()*X*A);
    //rowvec Ktemp = inv(B.t()*X*B + R) * B.t()*X*A;
    //Ktemp.print("Ktemp:");
    return Ktemp;
}

lqrResult_t lqrSteerControl(vec states, mat refData,
                            double pe, double pth_e, double L,
                            mat Q, mat R, double dt){
    minLatError_idx_t resMinLatError_idx = findMinLatError_idx(states, refData);
    double e = resMinLatError_idx.latError;
    uword index = resMinLatError_idx.index;

    double k = refData(index, 3);
    double th_e = toPi(states(2) - refData(index, 2));

    mat A = zeros<mat>(4, 4);
    A(0, 0) = 1; A(0, 1) = dt; A(1, 2) = states(3); A(2, 2) = 1;
    A(2, 3) = dt;

    vec B = zeros(4);
    B(3) = states(3)/L;

    rowvec K = dlqr(A, B, Q, R);
    vec x = zeros(4);
    x(0) = e; x(1) = (e-pe)/dt; x(2) = th_e; x(3) = (th_e-pth_e)/dt;

    double ff = atan2(L*k, 1); // 前馈转向
    double fb = toPi(dot(-K, x)); // 反馈转向
    double delta = 1*ff + 1*fb;

    lqrResult_t res;
    res.delta = delta;
    res.index = index;
    res.e = e;
    res.th_e = th_e;
    return res;
}

vec updateStates(vec states, double a, double delta,
                 double dt, double L, double maxSteer){
    //delta = max(min(maxSteer, delta), -maxSteer);

    if (delta >= maxSteer)
        delta = maxSteer;
    else if (delta <= -maxSteer)
        delta = - maxSteer;

    vec tempNewStates(4);
    tempNewStates[0] = states[0] + states[3]*cos(states[2])*dt;
    tempNewStates[1] = states[1] + states[3]*sin(states[2])*dt;
    tempNewStates[2] = states[2] + states[3]*dt*tan(delta)/L;
    tempNewStates[3] = states[3] + a*dt;
    return tempNewStates;
}

int main()
{
    // 总体跟踪效果：在参考路径点附近存在抖动
    double Kp = 0.8; // 过大不跟踪了
    double dt = 0.05;
    double L = 2.91;
    double targetSpeed = 72 * tompers; // m/s
    double maxSteer = 45 * deg2rad; // rad

    mat Q, R;
    Q.eye(4, 4);
    Q = 100*Q;
    R.eye(1, 1);
    R = 1*R;

    vector<double> XRef;
    vector<double> YRef;
    mat posRef = readPosPoints(XRef, YRef);

    // 计算一阶导数
    vec pd(posRef.n_rows);
    for(uword i = 0; i < posRef.n_rows-1; i++){
        pd(i) = (posRef(i+1, 1) - posRef(i, 1))/(posRef(i+1, 0) - posRef(i, 0));
    }
    pd(posRef.n_rows-1) = pd(posRef.n_rows-2);

    // 计算二阶导数
    vec pdd(posRef.n_rows);
    for(uword i = 1; i < posRef.n_rows-1; i++){
        pdd(i) = (posRef(i+1, 1) - 2*posRef(i, 1) + posRef(i-1, 1)) / pow(0.5*(posRef(i+1, 0) - posRef(i-1, 0)), 2);
    }
    pdd(0) = pdd(1);
    pdd(posRef.n_rows-1) = pdd(posRef.n_rows-2);

    // 计算曲率
    vec kappa(posRef.n_rows);
    for(uword i = 0; i < posRef.n_rows; i++){
        kappa(i) = pdd(i) / pow((1+pow(pd(i), 2)), 1.5);
    }

    // 打包参考数据
    vec refX = XRef;
    vec refY = YRef;
    vec refYaw = toPiVec(atan(pd));
    mat refData;
    refData = join_horiz(refX, refY, refYaw, kappa);

    double pe = 0;
    double pth_e = 0;

    // 初始Y值 过大 跟踪效果不佳
    vec initStates = {0, 0.1, 0, 0}; // X Y yaw v
    vec states = initStates;
    mat statesActual;
    statesActual = join_horiz(statesActual, initStates);

    uword index = 0;
    while(index < refX.n_elem-1){
        lqrResult_t resLQR = lqrSteerControl(states, refData,
                                             pe, pth_e, L, Q, R, dt);
        index = resLQR.index;
        pe = resLQR.e;
        pth_e = resLQR.th_e;

        /*
        if(abs(resLQR.e) > 4){
            cout << "The e is too large !!!" << endl;
            break;
        }
        */
        // 这里 除dt 相当于把Kp放大了
        double a = Kp * (targetSpeed - states(3))/dt;

        vec newStates =  updateStates(states, a, resLQR.delta,
                                     dt, L, maxSteer);
        states = newStates;

        // 保存实际的states
        statesActual = join_horiz(statesActual, newStates);
    }

    // 保存实际的位置点
    vector<double> XActual;
    vector<double> YActual;
    for(uword i = 0; i < statesActual.n_cols; i++){
        for(uword j = 0; j < statesActual.n_rows-2; j++){
            if(j == 0)
                XActual.push_back(statesActual(j, i));
            else
                YActual.push_back(statesActual(j, i));
        }
    }

    // 用于动态画出实际的点
       vector<double> plotXActual;
       vector<double> plotYActual;
       for(uword i = 0; i < XActual.size(); i++) {
           plotXActual.push_back(XActual[i]);
           plotYActual.push_back(YActual[i]);
           if (i % 10 == 0) {
               plt::clf();
               plt::named_plot("refPos", XRef, YRef, "r-");
               plt::named_plot("posActual", plotXActual, plotYActual, "b*");

               // Set x-axis to interval [0,1000000]
               plt::xlim(-10, 210);

               // Add graph title
               plt::title("LQR ErrorModel 1");

               // Enable legend.
               plt::legend();
               // Display plot continuously
               plt::pause(0.1);
           }
       }
       plt::show();

    return 0;
}
